//
// Created by kunnkali on 2023/5/23.
//

#ifndef TNET_TENSOR_H
#define TNET_TENSOR_H

#include "Option.h"
#include "Prefix.h"
#include "Result.h"
#include <memory>
#include <vector>

namespace TNet {

    enum TensorType { TensorTypeInput,
                      TensorTypeOutput,
                      TensorTypeUnknown };

    /**
     * @enum TensorDataType
     * This enum specifies supported TensorFlow Lite tensor data types.
     */
    enum TensorDataType {
        /** Tensor data type not available. This indicates an error with the model. */
        TensorDataTypeNoType = 0,

        /** 32-bit single precision floating point. */
        TensorDataTypeFloat32 = 1,

        /** 16-bit half precision floating point. */
        TensorDataTypeFloat16 = 2,

        /** 32-bit signed integer. */
        TensorDataTypeInt32 = 3,

        /** 8-bit unsigned integer. */
        TensorDataTypeUInt8 = 4,

        /** 64-bit signed integer. */
        TensorDataTypeInt64 = 5,

        /** Boolean. */
        TensorDataTypeBool = 6,

        /** 16-bit signed integer. */
        TensorDataTypeInt16 = 7,

        /** 8-bit signed integer. */
        TensorDataTypeInt8 = 8,

        /** 64-bit double precision floating point. */
        TensorDataTypeFloat64 = 9,
    };

    enum TensorDataFormat {
        NCHW    = 0,
        NHWC    = 1,
        NC4HW4  = 2,
        NHWC4   = 3,
        UNKNOWN = 4,
        MIN     = NCHW,
        MAX     = UNKNOWN
    };

    class TensorImpl;

    class TNET_API Tensor {
    public:
        Tensor(void *tensor = nullptr, TensorType type = TensorType::TensorTypeUnknown, FrameworkTypeTag frameworkType = FrameworkTypeTag::kTflite);

        Tensor(Tensor &&t) noexcept;
        ~Tensor();

        Tensor& operator=(Tensor&&);

        std::string name() const;

        TensorType type() const;

        TensorDataType dataType() const;

        TensorDataFormat dataFormat() const;

        std::vector<size_t> shape() const;

        size_t size() const;

        Base::Result resize(const std::vector<size_t> &shape);

        Base::Result copyToTensor(const void *data, size_t size);// 将数据拷贝到TME的Tensor中

        // 从TME的tensor中向外拷贝数据
        Base::Result copyFromTensor(void *data, size_t size) const;

        Base::Result copyToTensor(const Tensor &other);

        /// 获取数据存储的裸指针，可能存在内存布局问题，谨慎使用
        void *getRawDataPointer();

    private:
        FrameworkTypeTag _frameworkType;
        std::unique_ptr<TensorImpl> _impl;

        TNET_DECLARE_NON_COPYABLE(Tensor);
    };

}// namespace TNet

#endif// TNET_TENSOR_H
